import tensorflow as tf
import numpy as np
class TDetector:
    def __init__(self,network):
        frozen_graph = tf.GraphDef()
        with open(network, 'rb') as f:
            frozen_graph.ParseFromString(f.read())

        tf_config = tf.ConfigProto()
        tf_config.gpu_options.allow_growth = True
        self.__session = tf.Session(config=tf_config)
        tf.import_graph_def(frozen_graph, name='')
        self.__input = self.__session.graph.get_tensor_by_name('image_tensor:0')
        self.__scores = self.__session.graph.get_tensor_by_name('detection_scores:0')
        self.__boxes = self.__session.graph.get_tensor_by_name('detection_boxes:0')
        self.__classes = self.__session.graph.get_tensor_by_name('detection_classes:0')
        self.__num_detections = self.__session.graph.get_tensor_by_name('num_detections:0')
        self.count=0
    def detection(self,img):
        scores,boxes,classes,num_detections = self.__session.run((self.__scores,self.__boxes,self.__classes,self.__num_detections),feed_dict={self.__input:img[None,...]})
        scores,boxes,classes,num_detections=scores[0],boxes[0],classes[0],int(num_detections[0])
        h,w,_=img.shape
        self.scores,self.boxes,self.classes,self.count=scores[0:num_detections],boxes[0:num_detections],classes[0:num_detections],int(num_detections)
        for i in range(num_detections):
            self.boxes[i]=boxes[i]*np.array([h,w,h,w])  
    def __del__(self):
        self.__session.close()
